# pre packages
from myglobal import os, sys, LINES

# sys packages
from pylab import np
import obspy as ob
from glob import glob
from tqdm import tqdm
from obspy.core.utcdatetime import UTCDateTime
import pandas as pd
from scipy.stats import linregress
import datetime
from scipy import interpolate


# self packages
from utils.loc import load_loc, get_distance,sort_data_by_distance
from utils.math import measure_shift_fft,remove_point_skip
from utils.h5data import get_event_data, save_h5_data, read_h5_data,h5glob
from utils.plot import plot_scatters,plot_traces
from utils.trace import get_tapered_slices, safe_filter

# cmd
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-debug', action='store_true', help='method: DEBUG')
parser.add_argument('-figroot', default='figures/8.merge.dtshift', help='root to save figs')
parser.add_argument('-input', default='',  help='file pattern to load')
args = parser.parse_args()
print(args)

DEBUG= args.debug
INPUT = args.input
FIG_ROOT = args.figroot

datasets,args_infile= read_h5_data(INPUT,['all_groups','te'], group_name='metadata', read_attrs=True)
STATS = [i.decode() for i in datasets[0]]
te = datasets[1]
dte = te[1]-te[0]

print(args_infile)
DATE=args_infile['date']
EMARKER = args_infile['emarker']

FIG_ROOT = f'{FIG_ROOT}/8A.{DATE}.{EMARKER}'
if not os.path.exists(FIG_ROOT):
    os.makedirs(FIG_ROOT)

for name in STATS:
    
    data_keys = h5glob(INPUT, '*',group_name=f'{name}/merged',object_type='datasets', STRIP_GROUP_NAME=True)
    data_items = read_h5_data(INPUT, data_keys, group_name=f'{name}/merged')
    data = {data_keys[i]:data_items[i] for i in range(len(data_keys))}

    health_items= read_h5_data(INPUT, data_keys, group_name=f'{name}/health')
    health = {data_keys[i]:health_items[i] for i in range(len(data_keys))}

    phy_keys = h5glob(INPUT, '*',group_name='/phy',object_type='datasets', STRIP_GROUP_NAME=True)
    phy_items = read_h5_data(INPUT, phy_keys, group_name='/phy')
    data_phy = {phy_keys[i]:phy_items[i] for i in range(len(phy_keys))}

    from pylab import figure, plt
    from matplotlib import rcParams, gridspec
    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '8'
    DPI=600
    SCALE = 200/50  # m per ms
    plt.close('all')
    fig = figure(figsize=[5,10], dpi=DPI)
    NF = len(data_keys)+len(phy_keys)
    gs = fig.add_gridspec(NF,1,hspace=0)
    cmap = plt.get_cmap('tab10') 
    XLIM = [3,20]
    num_keys = len(data_keys)
    
    for i, key in enumerate(data_keys):
        ax = fig.add_subplot(gs[i,0])
        x = data[key]
        health_i = health[key]
        # x[np.abs(x)>50]=0
        # x = safe_filter(x, dte, ftype='bandpass', freqmin=0.1, freqmax=4)
        x = safe_filter(x, dte, ftype='lowpass', freq=3)
        x[health_i==0]=None
        color = cmap(i % cmap.N)  # 根据索引i选择颜色
        ax.plot(te,x, '-', color=color, lw=1, label=f'{key[50:]}')

        ax.set_xlim(XLIM)
        ax.set_ylim([-50,50])
        ax.set_ylabel('dt(ms)')
        ax.set_xlabel('')
        ax.set_xticklabels([])
        ax.legend(loc='lower right')

    for i, key in enumerate(phy_keys):
        # ax = fig.add_subplot(len(phy_keys)*2, 1, i+len(phy_keys)+1)
        ax = fig.add_subplot(gs[i+len(data_keys),0])
        x = data_phy[key]
        x = safe_filter(x, dte, ftype='lowpass', freq=3)
        ax.plot(te,x, '-', color='tab:red', lw=1, label=f'{key}')
        ax.scatter(te, x, s=10, c='gray', marker='^', edgecolors='none', label=None)

        ax.set_xlim(XLIM)
        # ax.set_ylim([-50,50])
        ax.set_ylabel('dt(ms)')
        ax.set_xlabel('')
        ax.set_xticklabels([])
        ax.legend(loc='lower right')

    fig.tight_layout()
    OUTFILE = f'{FIG_ROOT}/8A.{name}.{EMARKER}.png'
    print(OUTFILE)
    fig.savefig(OUTFILE)


# 绘制频率分析图

